
import wave
import contextlib
import numpy as np
import matplotlib.pyplot as plt
import math
import copy
import os
from scipy.io import wavfile
from pydub import AudioSegment
from numpy import set_printoptions

def get_wav_time(wav_path):
    '''
    获取音频文件是时长

    :param wav_path: 音频路径
    :return: 音频时长 (单位秒)
    '''
    with contextlib.closing(wave.open(wav_path, 'r')) as f:
        frames = f.getnframes()
    rate = f.getframerate()
    duration = frames / float(rate)
    return duration


def get_ms_part_wav(main_wav_path, start_time, end_time, part_wav_path):
    '''
    音频切片，获取部分音频 单位是毫秒级别

    :param main_wav_path: 原音频文件路径
    :param start_time:  截取的开始时间
    :param end_time:  截取的结束时间
    :param part_wav_path:  截取后的音频路径
    :return:
    '''
    start_time = int(start_time)
    end_time = int(end_time)
    sound = AudioSegment.from_mp3(main_wav_path)
    word = sound[start_time:end_time]
    word.export(part_wav_path, format="wav")

def get_wave_info(wavfile):
    f = wave.open(wavfile)
    SampleRate = f.getframerate()
    frames = f.getnframes()
    wav_time = frames / float(SampleRate)  # 单位为s

    return int(wav_time*1000)

def get_audio_list(data_path):
    audio_list=[]
    for file in os.listdir(data_path):
        a=os.path.splitext(file)
        if a[1]==".wav":
            audio_list.append(file)
    return audio_list

def split_audio_frame():
    dataset_path="/Users/leslie/Data/paper_audio/EMO_DB"
    processed_dir="/Users/leslie/Data/paper_audio/EMO_DB_processed"
    audio_list=get_audio_list(dataset_path)
    try:
        os.mkdir(processed_dir)
    except:
        pass
    #print(audio_list)
    from tqdm import tqdm
    for file in tqdm(audio_list):
        wav_file = file
        duration=get_wave_info(os.path.join(dataset_path,wav_file))
        frame_int=20
        sub_audio_start_ms=0
        sub_audio_end_ms=30
        storage_path=os.path.join(processed_dir,wav_file)
        try:
            os.mkdir(storage_path)
        except:
            pass
        while(sub_audio_end_ms<duration):

            sub_wav_name=os.path.splitext(wav_file)[0]+"_"+str(sub_audio_start_ms)\
                         +"_"+str(sub_audio_end_ms)+os.path.splitext(wav_file)[1]
            sub_audio_storage_path=os.path.join(storage_path,sub_wav_name)
            #print(sub_audio_storage_path)
            wav_path=os.path.join(dataset_path,wav_file)
            get_ms_part_wav(wav_path,sub_audio_start_ms,sub_audio_end_ms,sub_audio_storage_path)
            sub_audio_start_ms+=frame_int
            sub_audio_end_ms+=frame_int

if __name__ == '__main__':
    split_audio_frame()